import aiohttp
from abc import abstractmethod
from typing import Any, AsyncGenerator, Dict, List, Tuple

from evalscope.perf.arguments import Arguments
from evalscope.perf.utils.benchmark_util import BenchmarkData


class ApiPluginBase:

    def __init__(self, param: Arguments) -> None:
        self.param = param
        self.model_path = param.tokenizer_path

    @abstractmethod
    def build_request(self, messages: List[Dict], param: Arguments = None) -> Dict:
        """Build a api request body.

        Args:
            messages (List[Dict]): The messages generated by dataset.
            param (QueryParameters): The query parameters.

        Raises:
            NotImplementedError: Not implemented.

        Returns:
            Dict: The api request body.
        """
        raise NotImplementedError

    @abstractmethod
    def parse_responses(self, responses: List[Dict], request: str = None, **kwargs: Any) -> Tuple[int, int]:
        """Parser responses and return number of request and response tokens.

        Args:
            responses (List[Dict]): List of http response body, for stream output,
                there are multiple responses, each is bytes, for general only one.
            request (str): The json string of request.

        Returns:
            Tuple: (Number of prompt_tokens and number of completion_tokens).
        """
        raise NotImplementedError

    @abstractmethod
    async def process_request(
        self, client_session: aiohttp.ClientSession, url: str, headers: Dict, body: Dict
    ) -> BenchmarkData:
        """Process the HTTP request and handle the response.

        Args:
            client_session: The aiohttp client session
            url: The request URL
            headers: The request headers
            body: The request body

        Returns:
            BenchmarkData: The benchmark data including response and timing info.
        """
        raise NotImplementedError

    @staticmethod
    def replace_values(input_json: Any, model: str, prompt: str):
        if isinstance(input_json, dict):
            for key, value in input_json.items():
                if isinstance(value, str):
                    input_json[key] = value.replace('%m', model).replace('%p', prompt)
                else:
                    ApiPluginBase.replace_values(value, model, prompt)
        elif isinstance(input_json, list):
            for idx, item in enumerate(input_json):
                if isinstance(item, str):
                    input_json[idx] = item.replace('%m', model).replace('%p', prompt)
                else:
                    ApiPluginBase.replace_values(item, model, prompt)
        elif isinstance(input_json, str):
            input_json = input_json.replace('%m', model).replace('%p', prompt)
        else:
            pass
